package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"golang.org/x/tools/go/ast/astutil"
	"io/fs" // Go 1.16+ 用于 WalkDir 和文件权限
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	// "sync/atomic" // 在当前实现中，Load/Store 更直观，暂时不用 atomic
	// "unsafe" // 配合 atomic 使用，当前不需要
)

// ============================================================================
// 1. 追踪逻辑
// ============================================================================

var (
	// goroutineSpace 用于在 runtime.Stack 输出中识别goroutine ID的前缀
	goroutineSpace = []byte("goroutine ")
	// goroutineIndent 使用 sync.Map 存储每个 goroutine 的当前调用栈深度（用于缩进）
	// key: goroutine ID (uint64), value: indent level (int64)
	goroutineIndent sync.Map
)

// 定义ANSI转义序列颜色常量,用于控制台输出着色
const (
	colorReset   = "\033[0m"  // 重置所有格式
	colorGreen   = "\033[32m" // 绿色，用于 Goroutine ID
	colorBlue    = "\033[34m" // 蓝色，用于函数名
	colorMagenta = "\033[35m" // 品红色，用于行号
	// traceFuncName 是我们注入和追踪的函数的名字
	traceFuncName = "Trace" // 硬编码追踪函数名，简化单文件逻辑
)

// curGoroutineID 获取当前 goroutine 的 ID。
// 注意：此方法依赖 runtime.Stack 的输出格式，可能随 Go 版本变化，且有性能开销。
// 在生产环境中，可能需要更可靠的 goroutine 追踪方法。
func curGoroutineID() uint64 {
	var buf [64]byte                  // 分配一个小的栈缓冲区
	n := runtime.Stack(buf[:], false) // 获取当前 goroutine 的堆栈信息（非所有goroutine）
	// StSprintf is like fmt.Sprintf but writes into buf instead of allocating a new
	// string. It returns the number of bytes written.
	idField := bytes.TrimPrefix(buf[:n], goroutineSpace) // 移除 "goroutine " 前缀
	// Find the space separating the ID from the rest of the stack trace.
	i := bytes.IndexByte(idField, ' ')
	if i < 0 {
		// 如果格式不符合预期（例如，堆栈信息不完整）
		fmt.Fprintf(os.Stderr, "警告: 获取 goroutine ID 时无法解析堆栈: %q\n", buf[:n])
		return 0 // 返回 0 表示失败
	}
	// Parse the goroutine ID.
	id, err := strconv.ParseUint(string(idField[:i]), 10, 64)
	if err != nil {
		fmt.Fprintf(os.Stderr, "警告: 解析 goroutine ID '%s' 失败: %v\n", string(idField[:i]), err)
		return 0 // 返回 0 表示失败
	}
	return id
}

// printTrace 打印格式化的跟踪日志（函数进入或退出）
func printTrace(id uint64, name string, line int, arrow string, indent int) {
	// 生成缩进字符串
	indents := strings.Repeat("  ", indent) // 使用两个空格作为缩进单位
	// 格式化输出到标准输出
	fmt.Printf("%sgid[%05d]%s %s%s %s%s%s:%s%d%s\n",
		colorGreen, id, colorReset, // Goroutine ID (绿色)
		indents, arrow, // 缩进和箭头 ("->" 或 "<-")
		colorBlue, name, colorReset, // 函数名 (蓝色)
		colorMagenta, line, colorReset) // 行号 (品红色)
}

// Trace 是被注入到目标函数中的追踪函数。
// 它在函数进入时被调用，打印进入信息，并返回一个闭包函数。
// 这个闭包函数应该通过 defer 调用，在函数退出时打印退出信息。
func Trace() func() {
	// 获取调用者的程序计数器(pc)、文件名(file)、行号(line)
	pc, _, line, ok := runtime.Caller(1) // 参数 1 表示跳过 Trace 函数本身，获取调用者的信息
	if !ok {
		fmt.Fprintf(os.Stderr, "警告: runtime.Caller(1) 失败，无法获取调用者信息\n")
		return func() {} // 返回一个无操作的函数
	}

	// 根据程序计数器获取函数信息
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		fmt.Fprintf(os.Stderr, "警告: runtime.FuncForPC(%v) 失败，无法获取函数信息\n", pc)
		return func() {} // 返回一个无操作的函数
	}
	name := fn.Name()       // 获取函数的完整名称（如 "main.myFunction"）
	gid := curGoroutineID() // 获取当前 Goroutine 的 ID

	// --- 函数进入逻辑 ---
	// 原子地加载或存储当前 goroutine 的缩进级别。
	// 如果 gid 不存在，则存储 0 并返回 0；否则加载并返回现有值。
	rawIndent, _ := goroutineIndent.LoadOrStore(gid, int64(0))
	indent := rawIndent.(int64) // 类型断言为 int64

	// 打印函数进入的跟踪信息，使用当前缩进级别
	printTrace(gid, name, line, "->", int(indent))

	// 为下一次在此 goroutine 中的调用（或此函数的子调用）增加缩进级别
	// 注意：直接 Store 比原子加法更简单，并且对于缩进跟踪足够
	goroutineIndent.Store(gid, indent+1)

	// --- 返回函数退出时执行的闭包 ---
	return func() {
		// --- 函数退出逻辑 ---
		// 再次加载当前 goroutine 的缩进级别。
		// 注意：此时的值可能是进入时加 1 后的值，也可能被其他并发调用修改。
		currentRawIndent, loaded := goroutineIndent.Load(gid)
		if !loaded {
			// 理论上不应该发生，因为进入时至少存储了 0。
			fmt.Fprintf(os.Stderr, "警告: goroutine %d 的缩进信息在退出时丢失 (函数: %s)\n", gid, name)
			// 即使信息丢失，仍然尝试打印退出日志，使用进入时的缩进级别以保持配对感。
			printTrace(gid, name, line, "<-", int(indent))
			return
		}

		currentIndent := currentRawIndent.(int64) // 获取当前的缩进值

		// 打印函数退出的跟踪信息。
		// 为了使退出日志与对应的进入日志在视觉上对齐，我们使用 *进入时* 的缩进级别 `indent`。
		printTrace(gid, name, line, "<-", int(indent))

		// 函数已退出，减少此 goroutine 的缩进级别。
		newExitIndent := currentIndent - 1
		if newExitIndent <= 0 {
			// 如果缩进级别降至 0 或以下，从 map 中移除该 goroutine 的条目。
			// 这有助于防止 goroutine ID 重用时出现问题，并能释放少量内存。
			goroutineIndent.Delete(gid)
		} else {
			// 否则，更新缩进级别。
			goroutineIndent.Store(gid, newExitIndent)
		}
	}
}

// ============================================================================
// 2. AST 插桩逻辑
// ============================================================================

// instrumenter 结构体用于封装 AST 修改逻辑。
// 在单文件版本中，它不需要存储配置，因为追踪函数名是固定的。
type instrumenter struct{}

// NewInstrumenter 创建一个新的 instrumenter 实例。
func NewInstrumenter() *instrumenter {
	return &instrumenter{}
}

// Instrument 方法解析 Go 源文件，并向其中符合条件的函数注入 `defer Trace()()` 语句。
// filename: 要处理的 Go 源文件的路径。
// 返回值:
//   - []byte: 修改并格式化后的代码。
//   - error: 处理过程中发生的错误（如解析、格式化失败）。
//
// 如果文件无需修改（没有函数或所有函数已插桩），则返回 (nil, nil)。
func (a *instrumenter) Instrument(filename string) ([]byte, error) {
	// 1. 解析源文件
	fset := token.NewFileSet() // 创建文件集，用于管理源码位置信息
	// 解析文件，保留注释 (parser.ParseComments)
	curAST, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
	if err != nil {
		// 如果文件本身有语法错误，解析会失败
		return nil, fmt.Errorf("解析文件 %s 时出错: %w", filename, err)
	}

	// 2. 遍历和修改 AST
	var modified bool // 标记是否对 AST 进行了实际修改
	// 使用 astutil.Apply 高效地遍历和修改 AST 节点
	astutil.Apply(curAST, func(c *astutil.Cursor) bool {
		// 检查当前节点是否为函数声明 (*ast.FuncDecl)
		fd, ok := c.Node().(*ast.FuncDecl)
		// 如果不是函数声明，或者函数没有函数体 (例如接口方法、向前声明)，
		// 或者函数体为空，则不进行处理，继续遍历。
		if !ok || fd.Body == nil || len(fd.Body.List) == 0 {
			return true // 返回 true 表示继续遍历子节点及兄弟节点
		}

		// 尝试向函数体添加 defer 语句
		if a.addDeferStmt(fd) {
			modified = true // 如果成功添加了 defer 语句，标记为已修改
		}
		return true // 继续遍历
	}, nil) // post 函数为 nil，我们不需要在访问节点后执行操作

	// 3. 如果没有进行任何修改，则直接返回 nil, nil
	if !modified {
		return nil, nil
	}

	// 4. 格式化修改后的 AST
	// 在单文件版本中，Trace 函数就在 main 包内，不需要添加 import 语句。
	var buf bytes.Buffer // 创建缓冲区存储格式化后的代码
	// 使用 go/format 标准库将修改后的 AST 格式化回 Go 代码
	err = format.Node(&buf, fset, curAST)
	if err != nil {
		// 格式化通常不会失败，除非 AST 结构严重损坏
		return nil, fmt.Errorf("格式化 %s 的新代码时出错: %w", filename, err)
	}

	// 5. 返回格式化后的代码字节切片
	return buf.Bytes(), nil
}

// addDeferStmt 尝试向函数声明 (fd) 的函数体开头添加 "defer Trace()()" 语句。
// 此函数具有幂等性：如果目标 defer 语句已经存在，则不进行任何操作并返回 false。
// fd: 指向要修改的函数声明的指针 (*ast.FuncDecl)
// 返回值: 如果成功添加了 defer 语句，则返回 true；否则返回 false。
func (a *instrumenter) addDeferStmt(fd *ast.FuncDecl) (added bool) {
	// 基本检查：函数体和语句列表必须存在
	if fd.Body == nil || fd.Body.List == nil {
		return false
	}
	stmts := fd.Body.List

	// 检查 "defer Trace()()" 是否已经存在于函数体语句列表中
	// 这是为了防止重复注入
	for _, stmt := range stmts {
		// 检查语句是否为 defer 语句 (*ast.DeferStmt)
		ds, ok := stmt.(*ast.DeferStmt)
		if !ok {
			continue // 不是 defer 语句，检查下一个
		}

		// 检查 defer 调用的表达式是否为一个调用表达式 (*ast.CallExpr)，形如 defer f()()
		outerCall, ok := ds.Call.Fun.(*ast.CallExpr)
		if !ok {
			// defer 的目标不是一个立即调用，例如 defer f(a, b)，不匹配
			continue
		}

		// 检查这个内部调用表达式的函数部分是否为标识符 (*ast.Ident)
		ident, ok := outerCall.Fun.(*ast.Ident)
		if !ok {
			// 内部调用的不是一个简单的标识符，例如 defer pkg.Func()()，不匹配
			continue
		}

		// 检查标识符的名字是否是我们追踪函数的名字 "Trace"
		if ident.Name == traceFuncName {
			// 找到了完全匹配的 "defer Trace()()" 语句
			return false // 已经存在，无需添加，返回 false
		}
	}

	// --- 如果代码执行到这里，说明没有找到已存在的匹配 defer 语句 ---

	// 构造新的 defer 语句的 AST 节点: defer Trace()()
	newDeferStmt := &ast.DeferStmt{
		// DeferStmt 的 Call 字段代表 defer 后面跟着的表达式
		Call: &ast.CallExpr{ // 最外层的 ()
			// Fun 字段代表被调用的函数，这里是 Trace() 的结果
			Fun: &ast.CallExpr{ // 内层的 ()
				// Fun 字段代表被调用的函数，这里是标识符 Trace
				Fun: &ast.Ident{
					Name: traceFuncName, // "Trace"
				},
				// Args: 内层调用的参数列表，为空 []ast.Expr{}
				// Lparen, Rparen: 调用括号的位置信息，format.Node 会自动处理
			},
			// Args: 外层调用的参数列表，为空 []ast.Expr{}
		},
		// Defer: defer 关键字的位置信息，format.Node 会自动处理
	}

	// 将新的 defer 语句添加到函数体语句列表的开头
	// 创建一个新的切片，容量和长度比原来多 1
	newList := make([]ast.Stmt, 0, len(stmts)+1)
	// 添加新的 defer 语句作为第一个元素
	newList = append(newList, newDeferStmt)
	// 追加原来的所有语句
	newList = append(newList, stmts...)
	// 更新函数体的语句列表
	fd.Body.List = newList

	return true // 成功添加了 defer 语句，返回 true
}

// ============================================================================
// 3. 命令行处理逻辑
// ============================================================================

var (
	// write 标志决定是否将修改写回源文件，默认为 false (打印到 stdout)
	write = flag.Bool("w", false, "将结果写入源文件而不是标准输出")
)

// usage 打印命令行用法信息到标准错误
func usage() {
	fmt.Fprintf(os.Stderr, "用法: go run instrument.go [-w] [path]\n")
	fmt.Fprintf(os.Stderr, "       instrument [-w] [path] (编译后)\n\n")
	fmt.Fprintf(os.Stderr, "  path: 要处理的 Go 文件或目录。如果省略，默认为当前目录。\n")
	fmt.Fprintf(os.Stderr, "        会递归处理目录下的 .go 文件 (跳过 _test.go, vendor/, .*)。\n")
	fmt.Fprintf(os.Stderr, "选项:\n")
	flag.PrintDefaults() // 打印所有定义的标志的默认值和帮助信息
}

// main 函数是程序的入口点
func main() {
	flag.Usage = usage // 设置自定义的 Usage 函数
	flag.Parse()       // 解析命令行参数，填充 write 变量

	// 获取命令行传入的非标志参数（即路径）
	args := flag.Args()
	targetPath := "." // 默认目标路径为当前目录 "."
	if len(args) > 0 {
		targetPath = args[0] // 如果用户提供了路径，则使用该路径
	}

	// 获取目标路径的绝对路径，用于后续比较
	absTargetPath, err := filepath.Abs(targetPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "错误: 获取绝对路径 '%s' 失败: %v\n", targetPath, err)
		os.Exit(1)
	}

	// 检查目标路径是否存在
	targetInfo, err := os.Stat(absTargetPath)
	if err != nil {
		if os.IsNotExist(err) {
			fmt.Fprintf(os.Stderr, "错误: 路径 '%s' 不存在\n", absTargetPath)
		} else {
			fmt.Fprintf(os.Stderr, "错误: 访问路径 '%s' 时出错: %v\n", absTargetPath, err)
		}
		os.Exit(1) // 路径问题，直接退出
	}

	// 创建插桩器实例
	ins := NewInstrumenter()

	// --- 获取工具自身路径，用于跳过 ---
	selfPath := ""
	// 尝试获取编译后的可执行文件路径
	if execPath, execErr := os.Executable(); execErr == nil {
		selfPath, _ = filepath.Abs(execPath)
	}
	// 尝试获取通过 'go run' 运行时的源文件路径
	runFilePath := ""
	_, runFile, _, runOk := runtime.Caller(0)
	if runOk {
		runFilePath, _ = filepath.Abs(runFile)
	}
	// --- 工具路径获取结束 ---

	// 初始化处理结果计数器
	stats := struct {
		processed int // 已处理的文件数
		modified  int // 成功修改的文件数
		skipped   int // 跳过的文件/目录数
		errors    int // 处理失败的文件数
	}{}

	// --- 处理逻辑 ---
	if targetInfo.IsDir() {
		// 如果目标是目录，则递归遍历
		fmt.Printf("递归处理目录: %s\n", absTargetPath)
		// 使用 filepath.WalkDir (Go 1.16+) 进行遍历，它比 Walk 更高效且提供 DirEntry 信息
		err := filepath.WalkDir(absTargetPath, func(path string, d fs.DirEntry, walkErr error) error {
			// 1. 处理 WalkDir 本身可能遇到的错误 (例如权限问题)
			if walkErr != nil {
				fmt.Fprintf(os.Stderr, "警告: 访问路径 '%s' 出错: %v, 跳过\n", path, walkErr)
				stats.errors++
				// 可以选择返回 SkipDir 跳过整个子目录，但这里选择仅报告错误并继续
				return nil // 继续处理其他文件/目录
			}

			// 获取当前项的绝对路径
			absPath, pathErr := filepath.Abs(path)
			if pathErr != nil {
				fmt.Fprintf(os.Stderr, "警告: 获取绝对路径 '%s' 失败: %v, 跳过\n", path, pathErr)
				stats.errors++
				return nil
			}

			// 2. 跳过特定目录：vendor, 隐藏目录 (如 .git, .idea)
			if d.IsDir() {
				dirName := d.Name()
				if dirName == "vendor" || (dirName != "." && dirName != ".." && strings.HasPrefix(dirName, ".")) {
					fmt.Printf("跳过目录: %s\n", path)
					stats.skipped++
					return filepath.SkipDir // 指示 WalkDir 不要进入此目录
				}
				return nil // 是需要处理的目录，继续深入
			}

			// 3. 只处理文件
			if !d.Type().IsRegular() {
				fmt.Printf("跳过非文件类型: %s\n", path)
				stats.skipped++
				return nil
			}

			// 4. 只处理 .go 文件，且不是 _test.go 文件
			fileName := d.Name()
			if !strings.HasSuffix(fileName, ".go") || strings.HasSuffix(fileName, "_test.go") {
				if strings.HasSuffix(fileName, ".go") { // 是 _test.go
					fmt.Printf("跳过测试文件: %s\n", path)
					stats.skipped++
				} else { // 其他非 Go 文件
					// 可以选择不打印跳过非 Go 文件的信息，避免过多输出
					// fmt.Printf("跳过非 Go 文件: %s\n", path)
					// stats.skipped++
				}
				return nil // 跳过此文件
			}

			// 5. 跳过工具自身文件
			if absPath == selfPath || absPath == runFilePath {
				fmt.Printf("跳过工具自身文件: %s\n", path)
				stats.skipped++
				return nil
			}

			// --- 文件符合条件，进行处理 ---
			fmt.Printf("处理文件: %s\n", path)
			stats.processed++ // 增加已处理计数
			// 调用实际的文件处理函数
			modified, procErr := processFile(path, ins, *write)
			if procErr != nil {
				// 处理文件时发生错误（解析、格式化、写入等）
				fmt.Fprintf(os.Stderr, "  错误: 处理 %s 失败: %v\n", path, procErr)
				stats.errors++ // 增加错误计数
			} else if modified {
				// 文件成功修改
				writeMsg := ""
				if *write {
					writeMsg = "并写入"
				}
				fmt.Printf("  成功: 文件 %s 已插桩%s。\n", path, writeMsg)
				stats.modified++ // 增加修改计数
			} else {
				// 文件无需修改
				fmt.Printf("  信息: 文件 %s 无需修改。\n", path)
			}

			return nil // 继续遍历下一个文件/目录
		})

		if err != nil {
			// WalkDir 本身返回的最终错误（非 walkFn 中处理的错误）
			fmt.Fprintf(os.Stderr, "错误: 遍历目录 '%s' 时发生意外错误: %v\n", absTargetPath, err)
			stats.errors++
		}

	} else {
		// 如果目标是单个文件
		fileName := targetInfo.Name()
		absPath := absTargetPath // 单个文件时，absTargetPath 就是文件的绝对路径

		// 检查是否是需要处理的 Go 文件
		if !strings.HasSuffix(fileName, ".go") || strings.HasSuffix(fileName, "_test.go") {
			if strings.HasSuffix(fileName, "_test.go") {
				fmt.Printf("跳过测试文件: %s\n", absPath)
				stats.skipped++
			} else {
				fmt.Fprintf(os.Stderr, "错误: '%s' 不是有效的 Go 源文件。\n", absPath)
				stats.errors++
			}
		} else if absPath == selfPath || absPath == runFilePath {
			// 检查是否是工具自身
			fmt.Printf("跳过工具自身文件: %s\n", absPath)
			stats.skipped++
		} else {
			// 是需要处理的 Go 文件
			fmt.Printf("处理单个文件: %s\n", absPath)
			stats.processed++
			modified, procErr := processFile(absPath, ins, *write)
			if procErr != nil {
				fmt.Fprintf(os.Stderr, "  错误: 处理 %s 失败: %v\n", absPath, procErr)
				stats.errors++
			} else if modified {
				writeMsg := ""
				if *write {
					writeMsg = "并写入"
				}
				fmt.Printf("  成功: 文件 %s 已插桩%s。\n", absPath, writeMsg)
				stats.modified++
			} else {
				fmt.Printf("  信息: 文件 %s 无需修改。\n", absPath)
			}
		}
	}

	// --- 打印处理结果总结 ---
	fmt.Println("\n--- 处理完成 ---")
	fmt.Printf("已处理 Go 文件数: %d\n", stats.processed)
	fmt.Printf("成功修改文件数:  %d\n", stats.modified)
	fmt.Printf("跳过文件/目录数: %d\n", stats.skipped)
	fmt.Printf("处理失败文件数:  %d\n", stats.errors)

	// 如果处理过程中有任何错误，程序以非零状态码退出
	if stats.errors > 0 {
		os.Exit(1)
	}
}

// processFile 函数负责处理单个文件的插桩逻辑。
// path: 要处理的文件的绝对路径。
// ins: instrumenter 实例。
// writeToFile: 是否将修改写回文件。
// 返回值:
//   - bool: 表示文件是否被成功修改。
//   - error: 处理过程中发生的错误。
func processFile(path string, ins *instrumenter, writeToFile bool) (bool, error) {
	// 调用 Instrumenter 的核心方法进行代码分析和修改
	newSrc, err := ins.Instrument(path)
	if err != nil {
		// Instrument 内部发生错误（例如，源文件有语法错误导致解析失败）
		return false, fmt.Errorf("插桩分析失败: %w", err)
	}

	// 如果 Instrument 返回 nil，表示代码无需修改
	if newSrc == nil {
		return false, nil // 没有修改，也没有错误
	}

	// --- 代码已被修改 ---

	if !writeToFile {
		// 如果不需要写回文件 (-w 未指定)，只打印提示信息，避免大量代码输出
		fmt.Println("  (检测到修改，但未写入文件。使用 -w 选项写入)")
	} else {
		// 需要将修改后的代码写回源文件
		// 1. 获取原始文件信息，以便保留权限位
		originalInfo, statErr := os.Stat(path)
		perms := fs.FileMode(0644) // 默认使用 0644 权限
		if statErr == nil {
			perms = originalInfo.Mode().Perm() // 成功获取，则使用原始权限
		} else {
			// 获取原始权限失败，打印警告，继续使用默认权限
			fmt.Fprintf(os.Stderr, "警告: 无法获取文件 %s 的原始权限: %v, 将使用默认权限 0644\n", path, statErr)
		}

		// 2. 使用 os.WriteFile 将新内容写入文件，覆盖原有内容
		// WriteFile 会自动处理文件创建、截断和关闭
		err = os.WriteFile(path, newSrc, perms)
		if err != nil {
			// 写入文件时发生 IO 错误
			return true, fmt.Errorf("写入文件 '%s' 失败: %w", path, err) // 返回 true 表示已修改，但写入失败
		}
	}

	return true, nil // 文件已成功修改（无论是否写入）
}
